/*
 * @Author: wuqingchun
 * @Date: 2024-05-11 17:11:21
 * @Description: 文件锁的使用
 * @LastEditors: wuqingchun
 * @LastEditTime: 2024-05-21 20:00:48
 */
#include <thread>
#include <atomic>

#include <fcntl.h>
#include <unistd.h>

enum class FRLOCK_RESULT : uint8_t
{
    FRLOCK_SUCCESS,           // 加锁成功
    FRLOCK_SELF_LOCKED,       // 锁被自己持有
    FRLOCK_FAIL,              // 加锁失败
};

template <size_t FILE_RECORD_LOCK_COUNT>
class FileRecordLock
{
    static_assert(FILE_RECORD_LOCK_COUNT <= 1024, "The count of file record lock must be less than 1024 * 1024");
public:
    FileRecordLock()
    {
        for (uint32_t uiTmp = 0; uiTmp < FILE_RECORD_LOCK_COUNT; uiTmp++)
        {
            m_aui64RefCount[uiTmp].store(0);
        }
    }

    ~FileRecordLock()
    {

    }

    bool Init(const char* pcName)
    {
        m_iFd = open(pcName, O_RDWR | O_CREAT, 0666);
        if (m_iFd != -1)
        {
            int iData = 0;
            if (write(m_iFd, &iData, sizeof(iData)) != sizeof(iData))
            {
                close(m_iFd);
                m_iFd = -1;
            }
        }

        return m_iFd != -1;
    }

    FRLOCK_RESULT TryLock(uint32_t uiRecordIndex)
    {
        if (m_iFd == -1 || uiRecordIndex >= FILE_RECORD_LOCK_COUNT)
        {
            return FRLOCK_RESULT::FRLOCK_FAIL;
        }

        if (m_aui64RefCount[uiRecordIndex].fetch_add(1) > 0)
        {
            return FRLOCK_RESULT::FRLOCK_SELF_LOCKED;
        }

        if (__SetFileLock(F_WRLCK, uiRecordIndex) != 0)
        {
            m_aui64RefCount[uiRecordIndex].fetch_sub(1);
            return FRLOCK_RESULT::FRLOCK_FAIL; 
        }

        return FRLOCK_RESULT::FRLOCK_SUCCESS;
    }

    void UnLock(uint32_t uiRecordIndex)
    {
        if (m_iFd == -1 || uiRecordIndex >= FILE_RECORD_LOCK_COUNT)
        {
            return;
        }

        if (m_aui64RefCount[uiRecordIndex].fetch_sub(1) == 1)
        {
            (void)__SetFileLock(F_UNLCK, uiRecordIndex);
        }
    }

private:
    int __SetFileLock(short sLKType, uint32_t uiRecordIndex)
    {
        struct flock stFL = {  
                .l_type   = sLKType,    // 设置锁的类型：F_RDLCK（读锁）、F_WRLCK（写锁）或F_UNLCK（解锁）  
                .l_whence = SEEK_SET,   // 设置偏移量的起始位置  
                .l_start  = uiRecordIndex,   // 锁的开始位置  
                .l_len    = 1,          // 锁的长度，0表示到文件末尾  
                .l_pid    = 0,          // 锁的拥有者的PID，通常由系统填写  
        };

        return fcntl(m_iFd, F_SETLK, &stFL);
    }
private:
    int m_iFd = -1;
    std::atomic<uint64_t> m_aui64RefCount[FILE_RECORD_LOCK_COUNT]; /* 由于进程内针对同一个记录锁可以重复加锁，因此需要针对每一个记录维护一个引用计数 */
};


int main(int argc, char** argv)
{
    int iOffset = 0;
    if (argc > 1)
    {
        iOffset = atoi(argv[1]);
    }

    FileRecordLock<1024> objLock;
    if (!objLock.Init("test"))
    {
        return 0;
    }

    if (objLock.TryLock(iOffset) == FRLOCK_RESULT::FRLOCK_FAIL)
    {
        return 0;
    }

    if (objLock.TryLock(iOffset) == FRLOCK_RESULT::FRLOCK_FAIL)
    {
        return 0;
    }

    while (1)
    {
        printf("do something\n");
        std::this_thread::sleep_for(std::chrono::seconds(1));
    }

    return 0;
}